

#include <iostream>
#include <iomanip>
#include <fstream>
#include <string>
#include <sstream>
#include <stdexcept>
#include <vector>
#include <time.h>
#include <sys/time.h>
using std::cout;
using std::cerr;
using std::endl;
#include <cstdlib>

#include "gmrt_corr_lib.h"
#include "gmrt_corr_cuda.h"

#include "kernel_defs.hpp"
#include "maths_functions.hpp"
#include "device_interface.hpp"
#include "fill_device.hpp"
#include "unpack_device.hpp"
#include "phase_shift_device.hpp"
#include "fft_device.hpp"
#include "pfb_device.hpp"
#include "mac_device.hpp"
#include "stopwatch.hpp"
#include <cuda_runtime_api.h>
#include <cufft.h>
#include <math.h>
#define TIMESTR "%Y-%m-%d %T"

inline void start_profiling(const cudaEvent_t& start_event,
                            const cudaStream_t& stream=0)
{
#ifdef _PROFILE
  cudaEventRecord(start_event, stream);
#endif

}
inline float stop_profiling(const cudaEvent_t& start_event,
                            const cudaEvent_t& stop_event,
                            const cudaStream_t& stream=0)
{
  float elapsed = 0.f;
#ifdef _PROFILE
  cudaEventRecord(stop_event, stream);
  cudaEventSynchronize(stop_event);
  cudaEventElapsedTime(&elapsed, start_event, stop_event);
#endif
  return elapsed;
}


int gmrt_corr_cuda_create_events(gmrt_corr_profile_t * p)
{
	setDevice(0);   //function in device_interface GPU_CORR 7
  cudaEventCreate(&(p->start_event));
  cudaEventCreate(&(p->stop_event));
  p->h2d_time = 0;
  p->d2h_time = 0;
  p->unpack_time = 0;
  p->mac_time = 0;
  p->fft_time = 0;
  p->gpu_time = 0;
  p->phase_time = 0;
  p->run_count = 0;
  p->io_time = 0;
	return 0;
}

int gmrt_corr_cuda_alloc(gmrt_corr_data_t * d, gmrt_corr_params_t p)
{

 //functions in device_interface.hpp GPU_CORR 8
	size_t q,ant;
	for( q=0; q<p.max_queue_count; ++q ) {
    createQueue(d->queues[q]);
    mallocDevice(d->d_raw[q], p.max_raw_word_count * p.ant_count * sizeof(RawWord));
    mallocDevice(d->d_unpacked[q], p.max_unpacked_count*p.ant_count*sizeof(Real));    
  }
  mallocDevice(d->d_output, p.output_count*sizeof(Complex));
  mallocHost( d->h_output, p.output_count*sizeof(Complex));
 for(ant = 0;ant < p.ant_count;++ant)
 {
    mallocHost(  d->h_raw[ant], p.sample_count * sizeof(RawWord));
  }
	return 0;
}



int gmrt_corr_cuda_dealloc(gmrt_corr_data_t * d, gmrt_corr_params_t p)
{
	freeDevice(d->d_output);
  freeHost(d->h_output); 
  for( size_t q=0; q<p.max_queue_count; ++q ) {
    freeDevice(d->d_unpacked[q]); 
    freeDevice(d->d_raw[q]);
    destroyQueue(d->queues[q]);
  }
  for(size_t ant = 0; ant < p.ant_count; ++ant)
 {
    freeHost(d->h_raw[ant]);
 }
   
	return 0;
}

float gmrt_corr_cuda_run(gmrt_corr_data_t * d, gmrt_corr_params_t * p, gmrt_corr_profile_t * prof,
                         SourceParType *source, CorrType * corr)
{
  //initialization of variables... Most of them are needed for calculation of fringe and fstc values for phase shifting

  Complex zero;
  zero.x = 0.f;
  zero.y = 0.f;
  size_t ptr,ant;
  float time_update;
  cout << "bytes_per_sec = " << p->bytes_per_sec << endl;
  cout << "SAMPLES_PER_WORD = " << SAMPLES_PER_WORD << endl;
 time_update = float (p->freq_count*4.f/(p->bytes_per_sec));
  cout<< " freq_count = " << p->freq_count << endl;
  cout << "time_update = " << time_update << endl; 
  float fstc[p->ant_count];  
  float fringe_phase[p->ant_count];
  float fringe_rate[p->ant_count];
  float dd_t0[p->ant_count];
  size_t fft_count = 0;
  double total_time = 0;
  double total_bytes = 0;
  char buffer[30]; 
  for (size_t c = 0; c < p->ant_count; c++)
    {
      fstc[c] = 0.f + c * 0.f;
      fringe_phase[c] = 0.f;
      dd_t0[c] = 0.f ;       
    }

  //initialization of timers  

  Stopwatch timer;
  Stopwatch iotimer;
  Stopwatch fft_timer;
  Stopwatch gputimer;
  Stopwatch synctimer;
  

  timer.start();

  size_t bytes_to_read;
  size_t bytes_read;
  bool finished = false; // variable that controls the finish...If true the programs exits... Make this variable true in the while loop to
                         // exit after that iteration  
 
 
   FILE *fp = fopen("out.txt","ab");  //file to write output visibility into
   FILE *fl = fopen("log.txt","ab");  //file to write the fstc and fringe values into

   size_t queue_size;
   cufftResult error;
 
  // initializing the CUFFT plan and creating handles that can be used to perform FFT GPU_CORR 9

  cufftHandle fft_r2c_short_plan[p->max_queue_count];
   for (int i = 0 ; i < p->max_queue_count; i++)
{
  queue_size = math::min(p->max_queue_accum_size, p->gulp_accum_size-(i*p->max_queue_accum_size)); 

  error = cufftPlan1d(&(fft_r2c_short_plan[i]), p->freq_count, CUFFT_C2C, queue_size*p->ant_count);
  if( error != CUFFT_SUCCESS ) {
    throw std::runtime_error("Failed to create cufft 1d plan");
  }

   error = cufftSetStream(fft_r2c_short_plan[i], d->queues[i]);
  if( error != CUFFT_SUCCESS ) {
    throw std::runtime_error("Failed to set cufft stream");
  }

  cufftSetCompatibilityMode(fft_r2c_short_plan[i],CUFFT_COMPATIBILITY_NATIVE);


}


  // Begin run loop.. this loop performs streaming,unpacking,FFT,phase shifting,MAC and writes the output into out.txt. This loop exits when
  // finished variable becomes true GPU_CORR 10
     
  while( !finished ) {

     // calculating delay,fringe values and corresponding variables needed for their calculation
     fft_count = 0;
     total_time = total_bytes/p->bytes_per_sec;
     p->timestamp.tv_sec = (int) total_time;
     p->timestamp.tv_usec = (total_time - p->timestamp.tv_sec) * 1000000;
     p->timestamp.tv_sec += p->utc_start.tv_sec;
     fprintf(stdout,"  timestamps %ld   %ld\n",p->timestamp.tv_sec,p->timestamp.tv_usec);
     strftime(buffer,30,"%Y-%m-%d %T.",localtime(&(p->timestamp.tv_sec))); 

     calculate_delays(source, corr, p->timestamp, p->delays, p->fringes); // function that calculates fstc, fringe values and writes into
                                                                          // p->delays and p->fringes. Function in gmrt_delay_fns.c
     
    double int_samp;
    double samp_time=0.0025e-6;

     //calculating the actual fstc,fringe values from delay and fringe values obtained from 'calculate_delays' function...
     for (size_t c = 0;c < p->ant_count;c++)
     {
       int_samp = (p->delays[c].delay_t0/samp_time);  //no. of clocks
       //fstc[c] = (int_samp - (int)int_samp); // in seconds
        fstc[c] = 0.f; // in seconds
      // fringe_phase[c] = p->fringes[c].phase_ti;
       fringe_phase[c] = 0.f;
      //fringe_rate[c] =  p->fringes[c].dphase_t0;
       fringe_rate[c] =  0.f;
      //dd_t0[c] = p->delays[c].dd_t0/samp_time;       
       dd_t0[c] = 0.f;      
     
      fprintf(fl, "%s\t%s%6ld\t%15.12lf %25.12lf %15.12lf %15.8lf %15.8lf \n",corr->antenna[corr->sampler[c].ant_id].name,buffer,p->timestamp.tv_usec,fstc[c],p->delays[c].delay_t0/samp_time,p->delays[c].dd_t0,p->fringes[c].phase_ti,p->fringes[c].dphase_t0);
    }


   // reading the data files or shared memories in case of real-time test GPU_CORR 11
   for (ant = 0;ant < p->ant_count;++ant)
   {
    bytes_to_read = p->sample_count;
   //uncomment iotimer.start and iotimer.stop to get the timing for reading data files or shared memories
  
   // iotimer.start(); 

   bytes_read = d->io_function(d, ant, (char *) (d->h_raw[ant] ), bytes_to_read); // If reading from data files then function is in 
                                                                                   // gmrt_corr_lib.c 'gmrt_corr_file_read' function
                                                                                   // Incase of real-time test the function is in 
                                                                                   // psrdada/src/ipcio.c 'ipcio_read' function
    if (bytes_to_read != bytes_read)   
          {
            if (p->verbose)
              cerr << "EOD detected maybe??" << endl;
            finished = true;
          }
   // iotimer.stop();
         total_bytes += bytes_read;
     }


    prof->run_count++; // incrementing the iteration number... when this number matches the actual number of iterations given through command
                       // line the program exits
    if (p->verbose)
    {
      cout << "Beginning new accumulation..." << endl;
      cout << "=============================" << endl;
    }
       
    // Zero the accumulation result... filling zeroes in the memory corresponding to output visibility

     fill_device(d->d_output, d->d_output+p->output_count, zero); 

     // Iterate over accumulation gulps
     ptr = 0;
     size_t gulp_accum_size = p->gulp_accum_size;
  
     for( size_t accum=0; accum<p->science_accum_size; accum+=gulp_accum_size ) { // diving processing of one integration into number of
                                                                                  // small iterations depending upon gulp_size

      // Store gulp size, dealing with incomplete final gulp
      gulp_accum_size = math::min(gulp_accum_size, p->science_accum_size-accum);

      if (p->verbose)
      {
        cout << "Beginning new gulp accumulating " << gulp_accum_size << " cross correlations" << endl;
        cout << "---------------------------------------------------" << endl;
      }


      // Initializing the streams and their variables GPU_CORR 12       
      size_t queue_count = 0;
 
      // Iterate over queues in the gulp
      for( size_t queue=0; queue<gulp_accum_size; queue+=p->max_queue_accum_size ) {

        // Store queue's size, dealing with incomplete final queue
        d->queue_accum_size[queue_count] = math::min(p->max_queue_accum_size, gulp_accum_size-queue);

        // Calculate the unpack count and raw word count for this queue
        d->unpacked_count[queue_count] = d->queue_accum_size[queue_count] * 2*p->freq_count;
        d->raw_word_count[queue_count] = d->unpacked_count[queue_count] / SAMPLES_PER_WORD;
        if (p->verbose)
        {
          cout << "Using new queue accumulating " << d->queue_accum_size[queue_count] << " cross correlations" << endl;
          cout << "  => " << d->unpacked_count[queue_count] << " samples" << endl;
          cout << "  => " << d->raw_word_count[queue_count] << " raw words" << endl;
        }
        // Track the number of queues used for this gulp
        ++queue_count;       
      }
      if (p->verbose)
      {
        cout << "Total queues used in this gulp: " << queue_count << endl;
        cout << "Beginning data load..." << endl;
      }

       synchronize();  //Waiting for all threads to finish processing
 
       gputimer.start(); //starting gputimer... this calculates the total time for GPU procesing

       if(finished != true) {
 
       for( size_t q=0; q<queue_count; ++q ) { //Actual processing over streams

       if (p->verbose)
       cout << "Queue " << q << " copying data to device..." << endl;

       if(p->profiling)
       start_profiling(prof->start_event, d->queues[q]); //profiling H2D time


       // Copy data from host to device GPU_CORR 13
       
       for(size_t ant = 0; ant < p->ant_count; ++ant)
       {
        copyHostToDevice(d->d_raw[q] + (ant*d->raw_word_count[q]),  (d->h_raw[ant] + ptr) ,
                         d->raw_word_count[q]*sizeof(RawWord),
                         d->queues[q]); //function in device_interface.hpp
        }

        if(p->profiling)
        prof->h2d_time += stop_profiling(prof->start_event, prof->stop_event, d->queues[q]);
      
        if (p->verbose)
        cout << "Beginning compute..." << endl;
      
        // Compute
        if (p->verbose)
          cout << "Queue " << q << " unpacking data..." << endl;
 
       // Unpack
        if(p->profiling)
         start_profiling(prof->start_event, d->queues[q]); //profiling unpacking


        //unpacking the data... converting raw data into floats... function in unpack_device.hpp.. kernel name in unpack_device.hpp
        //                                                                                     'unpack_8bit_to_real_kernel' 
        unpack_to_real_device(d->d_raw[q], d->d_unpacked[q], 
                              d->raw_word_count[q]*p->ant_count,
                              sizeof(RawWord)*8/SAMPLES_PER_WORD,
                              d->queues[q]); 
        
        if(p->profiling)
          prof->unpack_time += stop_profiling(prof->start_event, prof->stop_event, d->queues[q]);
       
        if (p->verbose)
        {  
          cout << "Queue " << q << " FFTing data" << endl;
          cout << "  batch: " << d->queue_accum_size[q] * p->ant_count << endl;        
        }
        // FFT
        if(p->profiling)
         start_profiling(prof->start_event, d->queues[q]); //profiling FFT

        //FFT... function in fft_device.hpp... C2C FFT is performed and channels are rearranged using kernel 'finish_real_fft_kernel'
        fft_r2c_short_device(d->d_unpacked[q], (Complex*)d->d_unpacked[q],
                             2*p->freq_count,
                             d->queue_accum_size[q] * p->ant_count,
                             fft_r2c_short_plan,error, q, d->queues[q]);

        if(p->profiling)
         prof->fft_time += stop_profiling(prof->start_event, prof->stop_event, d->queues[q]);


        // Phase shift
        if(p->profiling)
         start_profiling(prof->start_event, d->queues[q]); //profiling phase shifting


        // phase shifting is done for all antennas from the fringe and fstc values calculated...Function in phase_shift_device.hpp... kernel
        // in 'phase_shift_kernel' GPU_CORR 14

        // Needs to be checked if calling a single kernel to perform phase shifting of all antennas improves performance and to see if lookup
        // table can be implemented for sine and cosine values

           for(size_t i=0;i<p->ant_count;i++)
       {
           phase_shift_device((Complex*)d->d_unpacked[q] + i*d->unpacked_count[q]/2 ,(Complex*)d->d_unpacked[q] + i*d->unpacked_count[q]/2,d->unpacked_count[q]/2 ,fstc[i],dd_t0[i], fringe_phase[i] ,fringe_rate[i] ,time_update ,p->freq_count, fft_count,d->queues[q]);
            }
        
  
        if(p->profiling)
          prof->phase_time += stop_profiling(prof->start_event, prof->stop_event, d->queues[q]);
     
     
        if (p->verbose)
          cout << "Queue " << q << " MACing data..." << endl;

        // MAC
        if(p->profiling)
         start_profiling(prof->start_event, d->queues[q]); //profiling MAC



        size_t in_stride1  = p->freq_count;
        size_t in_stride2  = d->unpacked_count[q] / 2; // Complex samples
        size_t out_stride1 = p->freq_count;
        size_t out_stride2 = p->ant_count * out_stride1;

        // MAC function in mac_device.hpp... kernel is 'mac_device_kernel' GPU_CORR 15   

        mac_device((Complex*)d->d_unpacked[q], in_stride1, in_stride2,
                   p->ant_count, p->freq_count, d->queue_accum_size[q],
                   d->d_output, out_stride1, out_stride2, p->do_full_polar,
                   d->queues[q]);


        if(p->profiling)
         prof->mac_time += stop_profiling(prof->start_event, prof->stop_event, d->queues[q]);


       fft_count += d->queue_accum_size[q];  
       ptr += d->raw_word_count[q];
      }
    }   // End of iteration over accumulation gulps
     synchronize();
     gputimer.stop(); //stop gputimer
  }


    if (p->verbose)
    {
      cout << "===========================" << endl;
      cout << "End of accumulation reached" << endl;
    }
 
    if (p->verbose) 
      cout << "Retrieving results from device..." << endl;
  
    // Copy back to host
        if(p->profiling)
        start_profiling(prof->start_event); //D2H profiling

       //Copy the result back to host... GPU_CORR 16
       copyDeviceToHost(d->h_output, d->d_output,p->output_count*sizeof(Complex)); 


        if(p->profiling)
        prof->d2h_time += stop_profiling(prof->start_event, prof->stop_event);
 
        
    if (p->verbose)
      cout << "Writing output to destination..." << endl;
    // Output
    synchronize(); // Wait for device->host copy to finish

    //writing results to the file out.txt
    if (p->write_output && !finished) {
      // HACK to avoid waiting ages when benchmarking
      if( p->ant_count <= 64 ) {

        size_t ij = 0;
        for( size_t i=0; i<p->ant_count; ++i ) {
          // Note: This is a little trick to divide by either 1 (single-pol)
          //       or 2 (full-polar) when skipping terms near the diagonal.
          size_t b = p->do_full_polar+1;
          for( size_t j=i/b*b; j<p->ant_count; ++j ) {
            ++ij;
            // Iterate over each frequency channel
            for( size_t f=0; f<p->freq_count; ++f ) {
              Complex cij = d->h_output[f+p->freq_count*(i+j*p->ant_count)];
              // Normalise the accumulations
              cij.x /= p->science_accum_size;
              cij.y /= p->science_accum_size;
              // Write the amplitude and phase
             fwrite(&cij,1,sizeof(cij),fp);
            }
          }
        }
      }
    }
   if(prof->run_count == p->integrations)
        finished = true;
  } // End of run loop

    fclose(fp);
    fclose(fl);
   for(size_t i=0;i<p->max_queue_count;++i)
  {
    cufftDestroy(fft_r2c_short_plan[i]); //Destroying the CUFFT plan
   }
  prof->io_time = iotimer.getTime() * 1000;
  prof->gpu_time = gputimer.getTime() * 1000;
  prof->gpu_sync_time = synctimer.getTime() * 1000;
  timer.stop();

  return timer.getTime();
}

